package org.example.strategy;

import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;

/**
 * 复合分片
 * @author shenguangyang
 */
public class CustomComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm<Long> {
    /**
     *
     * @param dataSourceNames 数据源集合
     *                        在分库时为所有分片库的集合 dataBaseNames, 比如表: product_order_0/product_order_1, 库ds0/ds1等
     *                        分表时为对应分片库中所有分片表的集合 tableNames
     * @param complexKeysShardingValue 分片属性, 包括
     *                                 loginTableName: 逻辑表
     *                                 columnNameAndShardingValuesMap: 存储多个分片键, 包括 key-value
     */
    @Override
    public Collection<String> doSharding(Collection<String> dataSourceNames, ComplexKeysShardingValue<Long> complexKeysShardingValue) {
        // 得到每个分片键对应的值
        Collection<Long> orderIdValues = this.getShardingValue(complexKeysShardingValue, "id");
        Collection<Long> userIdValues = this.getShardingValue(complexKeysShardingValue, "user_id");
        List<String> shardingSuffix = new ArrayList<>();
        // 对两个分片键取模组合方式, product_order_0_0、product_order_0_1、product_order_1_0、product_order_1_1
        for (Long userId : userIdValues) {
            for (Long orderId : orderIdValues) {
                String suffix = userId % 2 + "_" + orderId % 2;
                for (String dataSourceName : dataSourceNames) {
                    if (dataSourceName.equals(suffix)) {
                        shardingSuffix.add(dataSourceName);
                    }
                }
            }
        }
        return shardingSuffix;
    }

    /**
     * shardingValue: 分片属性, 包括
     * loginTableName: 逻辑表
     * columnNameAndShardingValuesMap: 存储多个分片键
     *  key: 分片key, id和user_id
     *  value: 分片value
     */
    private Collection<Long> getShardingValue(ComplexKeysShardingValue<Long> shardingValues, String key) {
        Collection<Long> valueSet = new ArrayList<>();
        Map<String, Collection<Long>> columnNameAndShardingValuesMap = shardingValues.getColumnNameAndShardingValuesMap();
        if (columnNameAndShardingValuesMap.containsKey(key)) {
            valueSet.addAll(columnNameAndShardingValuesMap.get(key));
        }
        return valueSet;
    }
}
